# -*- coding: utf-8 -*-
import pymysql
import pydot
import os
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from datetime import datetime
from config.db_config import DB_CONFIG  # 确保存在db_config.py文件


class ERGenerator:
    def __init__(self, config):
        self.config = config
        self.graph = pydot.Dot(graph_type='digraph', rankdir="LR",
                               bgcolor="#f8f9fa", fontname="PingFang SC")
        # MacOS环境路径配置（需通过Homebrew安装Graphviz）
        os.environ["PATH"] += os.pathsep + '/opt/homebrew/bin'  # Homebrew路径
        os.environ["PATH"] += os.pathsep + '/usr/local/bin'  # 传统路径

    def _get_metadata(self):
        """获取带注释的数据库元数据（含外键过滤）"""
        try:
            conn = pymysql.connect(**self.config)
            with conn.cursor() as cursor:
                cursor.execute("SHOW TABLES")
                tables = [table[0] for table in cursor.fetchall()]

                metadata = {}
                for table in tables:
                    # 获取字段元数据
                    cursor.execute(f"""
                        SELECT COLUMN_NAME, COLUMN_TYPE, COLUMN_KEY, COLUMN_COMMENT 
                        FROM INFORMATION_SCHEMA.COLUMNS 
                        WHERE TABLE_SCHEMA=DATABASE() 
                          AND TABLE_NAME='{table}'
                        ORDER BY ORDINAL_POSITION
                    """)
                    columns = [{
                        'name': row[0],
                        'type': row[1].decode('utf-8') if isinstance(row[1], bytes) else row[1],
                        'key': row[2],
                        'comment': row[3].decode('utf-8') if isinstance(row[3], bytes) else row[3]
                    } for row in cursor.fetchall()]

                    # 获取有效外键（过滤空值）
                    cursor.execute(f"""
                        SELECT COLUMN_NAME, REFERENCED_TABLE_NAME, REFERENCED_COLUMN_NAME
                        FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE 
                        WHERE TABLE_SCHEMA=DATABASE() 
                          AND TABLE_NAME='{table}'
                          AND REFERENCED_TABLE_NAME IS NOT NULL
                          AND REFERENCED_COLUMN_NAME IS NOT NULL
                    """)
                    foreign_keys = [{
                        'column': row[0],
                        'target_table': row[1],
                        'target_column': row[2]
                    } for row in cursor.fetchall() if all(row[1:])]  # 二次过滤

                    metadata[table] = {'columns': columns, 'foreign_keys': foreign_keys}
                return metadata
        except pymysql.err.OperationalError as e:
            print(f"❌ 数据库连接失败: {e}")
            print("💡 解决方案建议：")
            print("1. 确认DB_CONFIG中的端口号是否正确（默认3306）")
            print("2. 执行 sudo mysql -u root -p 登录后执行")
            print("   ALTER USER 'root'@'localhost' IDENTIFIED WITH mysql_native_password BY '新密码';")
            raise
        finally:
            if 'conn' in locals():
                conn.close()

    def _build_nodes(self, metadata):
        """构建带注释的ER图节点（支持中文）"""
        for table, data in metadata.items():
            # HTML表格结构（优化中文显示）
            label = '''<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="4">
                        <TR><TD COLSPAN="2" BGCOLOR="#e9ecef"><B>{}</B></TD></TR>'''.format(table)

            for col in data['columns']:
                # 主键特殊标识
                bg_color = "#fff3cd" if col['key'] == 'PRI' else "#ffffff"
                icon = "🔑" if col['key'] == 'PRI' else "→" if col['key'] == 'MUL' else " "

                label += '''<TR>
                              <TD BGCOLOR="{}" ALIGN="CENTER">{}</TD>
                              <TD ALIGN="LEFT"><FONT POINT-SIZE="12">{}</FONT><BR/>
                                <FONT COLOR="#6c757d" POINT-SIZE="10">{}</FONT></TD>
                            </TR>'''.format(
                    bg_color,
                    icon,
                    f"{col['name']} : {col['type']}",
                    col['comment'] or '无注释'
                )
            label += "</TABLE>>"

            node = pydot.Node(
                table,
                shape="none",
                label=label,
                fontname="PingFang SC"
            )
            self.graph.add_node(node)

    def _build_edges(self, metadata):
        """构建外键关系连线（含箭头样式优化）"""
        for table, data in metadata.items():
            for fk in data['foreign_keys']:
                # 检查目标表是否存在
                if fk['target_table'] not in metadata:
                    print(f"⚠️ 警告：目标表 {fk['target_table']} 不存在，跳过该外键")
                    continue

                edge = pydot.Edge(
                    table, fk['target_table'],
                    label=f" {fk['column']} → {fk['target_column']} ",
                    color="#4a8db7",
                    penwidth=1.5,
                    arrowhead="crow",  # 使用crow箭头表示外键
                    fontname="PingFang SC",
                    fontsize=10,
                    labelfontcolor="#6c757d"
                )
                self.graph.add_edge(edge)

    def generate(self, output_path="er_diagram.png"):
        """生成并展示ER图"""
        metadata = self._get_metadata()
        self._build_nodes(metadata)
        self._build_edges(metadata)

        # 或者直接这样写
        self.graph.write_png(output_path)

        # 用Matplotlib展示（适配Retina显示屏）
        plt.rcParams['font.sans-serif'] = ['PingFang SC']  # 设置中文字体
        img = mpimg.imread(output_path)
        plt.figure(figsize=(24, 18), dpi=100)
        plt.imshow(img)
        plt.axis('off')
        plt.title(f"数据库ER图 - {datetime.now().strftime('%Y-%m-%d %H:%M')}",
                  fontsize=14, pad=20)
        plt.tight_layout()
        plt.show()


if __name__ == "__main__":
    generator = ERGenerator(DB_CONFIG)
    generator.generate(output_path="../output/macOS_ER_diagram.png")